from Bio import Entrez
import os
import sys
import google.generativeai as genai
import time
import random
from datetime import datetime
from dotenv import load_dotenv  # 新增安全模块
# 安全加载环境变量（优先从非版本控制的.env文件读取）
def safe_load_config():
    """安全加载环境变量配置
    
    【新手必读】使用前需完成以下步骤：
    1. 注册获取API密钥：
       - NCBI账户：访问 https://www.ncbi.nlm.nih.gov/account/ 获取ENTREZ_API_KEY
       - Google AI：访问 https://makersuite.google.com/ 获取GENAI_API_KEY
    2. 在项目根目录创建.env文件（模板见下方）
    3. 填写真实密钥（勿上传到Git仓库！）
    4. 遇到问题请联系微信彭老师 📱13084473447
    
    .env文件模板内容：
    --------------------------------------------------
    ENTREZ_EMAIL=你的注册邮箱
    ENTREZ_API_KEY=你的NCBI_API密钥
    GENAI_API_KEY=你的Google_AI密钥
    --------------------------------------------------
    """
    # 定位.env文件路径（与当前文件同目录）
    env_path = os.path.join(os.path.dirname(__file__), '.env')
    
    # 检查配置文件是否存在
    if os.path.exists(env_path):
        load_dotenv(env_path)  # 加载环境变量
    else:
        raise FileNotFoundError("❌ 找不到.env文件，请先创建配置文件（见上方注释说明）")
    
    # 验证必需配置项
    required_keys = ['ENTREZ_EMAIL', 'ENTREZ_API_KEY', 'GENAI_API_KEY']
    missing = [key for key in required_keys if not os.getenv(key)]
    
    if missing:
        error_msg = f"⚠️ 缺少必需配置: {', '.join(missing)}\n  请检查.env文件或联系彭老师获取帮助"
        raise ValueError(error_msg)
    
    # 返回配置字典（开发者无需修改此部分）
    return {
        'email': os.environ['ENTREZ_EMAIL'],
        'entrez_key': os.environ['ENTREZ_API_KEY'],
        'genai_key': os.environ['GENAI_API_KEY']
    }
# 获取应用程序路径（添加路径安全检查）
def get_application_path():
    """获取安全的应用程序路径"""
    base_path = os.path.dirname(sys.executable) if getattr(sys, 'frozen', False) \
        else os.path.dirname(os.path.abspath(__file__))
    
    # 防止路径遍历攻击
    if not os.path.exists(base_path):
        raise FileNotFoundError(f"无效的路径配置: {base_path}")
    
    return base_path
# 初始化配置
try:
    config = safe_load_config()
    Entrez.email = config['email']          # 从环境变量读取
    Entrez.api_key = config['entrez_key']    # 从环境变量读取
    genai.configure(api_key=config['genai_key'])  # 安全配置GenAI
except (FileNotFoundError, ValueError) as e:
    print(f"配置错误: {str(e)}")
    sys.exit(1)


def search_pubmed(query_term, max_results=10):
    """
    使用 PubMed API 搜索文献，并返回文献 ID 列表。

    Args:
        query_term: 搜索关键词。
        max_results: 最大返回结果数量。

    Returns:
        文献 ID 列表。
    """
    try:
        handle = Entrez.esearch(db="pubmed", term=query_term, retmax=max_results)
        record = Entrez.read(handle)
        handle.close()
        return record["IdList"]
    except Exception as e:
        print(f"搜索PubMed时出错: {e}")
        return []


def get_paper_details(pmid):
    """
    获取论文的详细信息，包括标题、作者、期刊、发表日期等。
    
    Args:
        pmid: PubMed ID
        
    Returns:
        包含论文详细信息的字典
    """
    try:
        handle = Entrez.efetch(db="pubmed", id=pmid, retmode="xml")
        record = Entrez.read(handle)
        handle.close()
        
        article = record['PubmedArticle'][0]['MedlineCitation']['Article']
        journal = article['Journal']
        
        # 获取标题
        title = article['ArticleTitle']
        
        # 获取作者 (如果有)
        authors = []
        if 'AuthorList' in article:
            for author in article['AuthorList']:
                if 'LastName' in author and 'ForeName' in author:
                    authors.append(f"{author['LastName']} {author['ForeName']}")
                elif 'LastName' in author and 'Initials' in author:
                    authors.append(f"{author['LastName']} {author['Initials']}")
                elif 'CollectiveName' in author:
                    authors.append(author['CollectiveName'])
        
        # 获取期刊信息
        journal_name = journal['Title'] if 'Title' in journal else "Unknown Journal"
        
        # 获取发表日期
        pub_date = ""
        if 'PubDate' in journal['JournalIssue']:
            date_parts = journal['JournalIssue']['PubDate']
            if 'Year' in date_parts:
                pub_date = date_parts['Year']
                if 'Month' in date_parts:
                    pub_date = f"{date_parts['Month']} {pub_date}"
                if 'Day' in date_parts:
                    pub_date = f"{date_parts['Day']} {pub_date}"
        
        # 获取DOI (如果有)
        doi = ""
        if 'ELocationID' in article:
            for id in article['ELocationID']:
                if id.attributes['EIdType'] == 'doi':
                    doi = id
        
        # 获取摘要
        abstract = ""
        if 'Abstract' in article and 'AbstractText' in article['Abstract']:
            if isinstance(article['Abstract']['AbstractText'], list):
                for text in article['Abstract']['AbstractText']:
                    if isinstance(text, str):
                        abstract += text + " "
                    else:
                        abstract += text.title() + ": " + text + " "
            else:
                abstract = article['Abstract']['AbstractText']
        
        return {
            "pmid": pmid,
            "title": title,
            "authors": ", ".join(authors),
            "journal": journal_name,
            "pub_date": pub_date,
            "doi": doi,
            "abstract": abstract
        }
    
    except Exception as e:
        print(f"Error getting details for PMID {pmid}: {e}")
        return {
            "pmid": pmid,
            "title": "Error retrieving information",
            "authors": "",
            "journal": "",
            "pub_date": "",
            "doi": "",
            "abstract": ""
        }


def summarize_pubmed_abstract(pmid, gemini_api_key, query_topic):
    """
    使用 Gemini 模型总结 PubMed 文章摘要。

    Args:
        pmid: PubMed ID (例如：'35854100').
        gemini_api_key: Gemini API 密钥
        query_topic: 检索主题，用于提示模型关注相关内容

    Returns:
        总结后的摘要文本，如果发生错误则返回 None。
    """
    try:
        # 从 PubMed 获取文章信息
        paper_details = get_paper_details(pmid)
        abstract = paper_details["abstract"]
        
        if not abstract:
            return {
                "pmid": pmid,
                "title": paper_details["title"],
                "authors": paper_details["authors"],
                "journal": paper_details["journal"],
                "pub_date": paper_details["pub_date"],
                "summary": "未找到摘要。"
            }

        # 使用 Gemini 模型进行总结
        genai.configure(api_key=gemini_api_key)
        model = genai.GenerativeModel('gemini-1.5-pro')
        prompt = f"""请使用中文总结以下关于"{query_topic}"的PubMed文章摘要。
        包括关键发现、方法和意义，大约3-4个句子:
        
        标题: {paper_details['title']}
        摘要: {abstract}"""
        
        # 添加重试机制
        max_retries = 3
        for attempt in range(max_retries):
            try:
                response = model.generate_content(prompt)
                summary = response.text.strip()
                return {
                    "pmid": pmid,
                    "title": paper_details["title"],
                    "authors": paper_details["authors"],
                    "journal": paper_details["journal"],
                    "pub_date": paper_details["pub_date"],
                    "summary": summary
                }
            except Exception as e:
                if "RATE_LIMIT_EXCEEDED" in str(e) and attempt < max_retries - 1:
                    wait_time = (2 ** attempt) + random.random()  # 指数退避策略
                    print(f"Rate limit exceeded. Waiting {wait_time:.2f} seconds before retry...")
                    time.sleep(wait_time)
                else:
                    raise  # 重新抛出异常

    except Exception as e:
        print(f"An error occurred: {e}")
        return {
            "pmid": pmid,
            "title": paper_details.get("title", "Unknown"),
            "authors": paper_details.get("authors", ""),
            "journal": paper_details.get("journal", ""),
            "pub_date": paper_details.get("pub_date", ""),
            "summary": f"无法生成摘要: {str(e)}\n\n原始摘要: {paper_details.get('abstract', '')[:300]}..."
        }


def fetch_abstracts(id_list):
    """
    根据文献 ID 列表，获取文献摘要。

    Args:
        id_list: 文献 ID 列表。

    Returns:
        包含摘要的字符串，每篇摘要之间用换行符分隔。
    """
    if not id_list:
        return ""

    try:
        handle = Entrez.efetch(db="pubmed", id=id_list, idtype="pmid", rettype="abstract", retmode="text") # 明确指定 idtype="pmid"
        abstracts = handle.read()
        handle.close()
        return abstracts
    except Exception as e:
        print(f"获取摘要时出错: {e}")
        return ""


def summarize_abstracts_with_gemini(abstracts, gemini_api_key, query_topic):
    """
    使用 Gemini 模型，总结多篇文献摘要并生成综合总结。

    Args:
        abstracts: 包含摘要的字符串。
        gemini_api_key: 您的 Gemini API 密钥。
        query_topic: 检索主题

    Returns:
        总结结果字符串。
    """
    try:
        genai.configure(api_key=gemini_api_key)
        model = genai.GenerativeModel('gemini-2.0-flash')  # 使用最新的模型名称

        prompt_content = f"""请用中文总结以下关于"{query_topic}"的医学文献摘要。
        分析共同的发现、方法差异、临床意义以及可能的未来研究方向，提炼出最关键的信息:
        
        {abstracts}
        
        请提供一份全面但简明的总结，包括以下方面:
        1. 研究背景和重要性
        2. 主要发现和结论
        3. 临床实践建议
        4. 未来研究方向
        """

        # 添加重试机制
        max_retries = 5
        for attempt in range(max_retries):
            try:
                response = model.generate_content(prompt_content)
                summary = response.text.strip()
                return summary
            except Exception as e:
                if "RATE_LIMIT_EXCEEDED" in str(e) and attempt < max_retries - 1 or "quota" in str(e).lower() or "exhausted" in str(e).lower() or "429" in str(e):
                    wait_time = (2 ** attempt) + random.random()  # 指数退避策略
                    print(f"Rate limit exceeded. Waiting {wait_time:.2f} seconds before retry...")
                    time.sleep(wait_time)
                else:
                    raise  # 重新抛出异常

    except ImportError:
        return "Error: google-generativeai 库未安装。请先安装该库 (pip install google-generativeai) 并配置 Gemini API 密钥。"
    except Exception as e:
        return f"调用 Gemini API 出错: {e}。\n\n由于API限制，无法生成综合总结，但您仍可以阅读各文献的原始摘要。"


def process_individual_papers_with_quota_management(pubmed_ids, gemini_api_key, query_topic, batch_size=3):
    """
    分批处理论文并处理配额限制
    
    Args:
        pubmed_ids: PubMed ID列表
        gemini_api_key: Gemini API密钥
        query_topic: 检索主题
        batch_size: 每批处理的文献数量
    
    Returns:
        包含每篇论文摘要的列表
    """
    paper_summaries = []
    
    # 分批处理
    for i in range(0, len(pubmed_ids), batch_size):
        batch_ids = pubmed_ids[i:i+batch_size]
        print(f"\n处理批次 {i//batch_size + 1}/{(len(pubmed_ids) + batch_size - 1)//batch_size}...")
        
        batch_summaries = []
        for j, pmid in enumerate(batch_ids):
            print(f"  处理论文 {i+j+1}/{len(pubmed_ids)}，PMID: {pmid}...")
            
            try:
                summary = summarize_pubmed_abstract(pmid, gemini_api_key, query_topic)
                batch_summaries.append(summary)
                
                # 打印部分结果
                print(f"  ✓ 成功处理论文 {i+j+1}: {summary['title'][:50]}...\n")
                
                # 添加延迟，避免触发速率限制
                if j < len(batch_ids) - 1:
                    delay = random.uniform(5, 8)  # 随机延迟5-8秒
                    print(f"    等待 {delay:.2f} 秒...")
                    time.sleep(delay)
                    
            except Exception as e:
                error_msg = str(e)
                print(f"  ✗ 处理论文 PMID: {pmid} 时出错: {error_msg}")
                
                # 检查是否是配额限制错误
                if "quota" in error_msg.lower() or "exhausted" in error_msg.lower() or "429" in error_msg:
                    print("\n⚠️ API配额已耗尽，等待较长时间恢复...")
                    
                    # 获取论文基本信息，不使用AI总结
                    try:
                        paper_details = get_paper_details(pmid)
                        batch_summaries.append({
                            "pmid": pmid,
                            "title": paper_details.get("title", "未知标题"),
                            "authors": paper_details.get("authors", ""),
                            "journal": paper_details.get("journal", ""),
                            "pub_date": paper_details.get("pub_date", ""),
                            "summary": "API配额限制：无法生成摘要。原始摘要：\n" + paper_details.get("abstract", "无法获取原始摘要")[:300] + "..."
                        })
                    except:
                        batch_summaries.append({
                            "pmid": pmid,
                            "title": "API配额已耗尽",
                            "authors": "",
                            "journal": "",
                            "pub_date": "",
                            "summary": "由于API配额限制，无法生成摘要。请稍后再试或减少请求数量。"
                        })
                    
                    # 延长等待时间，希望配额会恢复
                    recovery_wait = 60  # 等待60秒
                    print(f"    等待 {recovery_wait} 秒后继续...")
                    time.sleep(recovery_wait)
                else:
                    batch_summaries.append({
                        "pmid": pmid,
                        "title": "处理出错",
                        "authors": "",
                        "journal": "",
                        "pub_date": "",
                        "summary": f"无法生成摘要: {error_msg}"
                    })
        
        # 保存这批处理的结果
        paper_summaries.extend(batch_summaries)
        
        # 批次之间等待较长时间
        if i + batch_size < len(pubmed_ids):
            batch_wait = random.uniform(15, 20)  # 批次间等待15-20秒
            print(f"\n批次处理完成，等待 {batch_wait:.2f} 秒后处理下一批...")
            time.sleep(batch_wait)
    
    return paper_summaries


def create_literature_report(query, pubmed_ids, paper_summaries, overall_summary, output_file=None):
    """
    创建文献报告并保存为文本文件。
    
    Args:
        query: 搜索查询
        pubmed_ids: PubMed ID列表
        paper_summaries: 每篇论文的摘要
        overall_summary: 总体摘要
        output_file: 输出文件名（可选）
    
    Returns:
        报告文本内容
    """
    # 如果没有提供输出文件名，生成一个基于日期和查询的文件名
    if output_file is None:
        date_str = datetime.now().strftime("%Y%m%d_%H%M%S")
        sanitized_query = "".join(c if c.isalnum() else "_" for c in query)[:30]
        output_file = os.path.join(get_application_path(), f"PubMed_Report_{sanitized_query}_{date_str}.txt")
    
    # 创建报告内容
    report = []
    
    # 添加标题和搜索信息
    report.append("=" * 80)
    report.append(f"医学文献综述报告: {query}")
    report.append("=" * 80)
    report.append(f"搜索日期: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    report.append(f"搜索关键词: {query}")
    report.append(f"文献数量: {len(pubmed_ids)}")
    report.append("")
    
    # 添加总体摘要
    report.append("=" * 80)
    report.append("综合总结")
    report.append("=" * 80)
    report.append(overall_summary)
    report.append("\n")
    
    # 添加每篇文献的摘要
    report.append("=" * 80)
    report.append("各篇文献摘要")
    report.append("=" * 80)
    
    for i, paper in enumerate(paper_summaries):
        report.append(f"[{i+1}] {paper['title']}")
        report.append(f"作者: {paper['authors']}")
        report.append(f"期刊: {paper['journal']}")
        report.append(f"发表日期: {paper['pub_date']}")
        report.append(f"PubMed ID: {paper['pmid']}")
        report.append("")
        report.append("摘要:")
        report.append(paper['summary'])
        report.append("")
        report.append("-" * 40)
    
    # 添加引用信息
    report.append("=" * 80)
    report.append("引用信息")
    report.append("=" * 80)
    for i, paper in enumerate(paper_summaries):
        report.append(f"[{i+1}] {paper['authors']}. {paper['title']}. {paper['journal']}. {paper['pub_date']}. PMID: {paper['pmid']}")
    
    # 添加生成说明
    report.append("")
    report.append("=" * 80)
    report.append("本报告由 PubMed 文献检索系统自动生成")
    report.append(f"生成日期：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    report.append("=" * 80)
    
    # 将报告内容写入文件
    report_text = "\n".join(report)
    try:
        with open(output_file, "w", encoding="utf-8") as f:
            f.write(report_text)
        
        print(f"报告已保存至: {os.path.abspath(output_file)}")
    except Exception as e:
        print(f"保存报告时出错: {e}")
        # 尝试使用备用路径
        try:
            desktop_path = os.path.join(os.path.expanduser("~"), "Desktop")
            backup_file = os.path.join(desktop_path, f"PubMed_Report_{date_str}.txt")
            with open(backup_file, "w", encoding="utf-8") as f:
                f.write(report_text)
            print(f"报告已保存至桌面: {backup_file}")
        except:
            print("无法保存报告。请检查权限或磁盘空间。")
    
    return report_text


def get_user_input():
    """
    获取用户输入的搜索关键词和文献数量
    
    Returns:
        搜索关键词和文献数量
    """
    print("\n" + "=" * 60)
    print(" " * 15 + "PubMed 文献检索与总结系统")
    print("=" * 60)
    
    # 获取搜索关键词
    while True:
        search_query = input("\n请输入您想检索的医学主题关键词 (建议使用英文): ").strip()
        if search_query:
            break
        print("关键词不能为空，请重新输入！")
    
    # 获取文献数量
    while True:
        try:
            paper_count = input("\n请输入您想检索的文献数量 (1-20): ").strip()
            paper_count = int(paper_count)
            if 1 <= paper_count <= 20:
                break
            elif paper_count > 20:
                print("警告：太花钱了，省着用！请输入不超过20的数字。")
            else:
                print("请输入一个介于1和20之间的整数！")
        except ValueError:
            print("请输入一个有效的整数！")
    
    # 确认信息
    print("\n" + "-" * 60)
    print(f"您将检索关于「{search_query}」的 {paper_count} 篇最新文献")
    print("-" * 60)
    
    return search_query, paper_count


if __name__ == "__main__":
    try:
        # 显示欢迎信息（保持不变）
        print("\n" + "=" * 70)
        print(" " * 20 + "PubMed 文献检索与总结系统")
        print(" " * 15 + "自动检索并总结医学研究文献的最新进展")
        print("=" * 70)
        
        # ================== 修改点1：API密钥获取方式 ==================
        your_gemini_api_key = input("\n请输入您的Gemini API密钥（必需）: ").strip()
        while not your_gemini_api_key:
            print("\033[91m[安全警告] 必须提供有效的API密钥！\033[0m")
            print(">>> 密钥获取方式：")
            print("1. 访问 https://ai.google.dev/ 注册账号")
            print("2. 在控制台创建API密钥（你可以暂时使用AIzaSyC7w6R6eNNn3Jo6C1Vbllx9WwFN_0FVsQE），但是不保证未来长期能用")
            print("3. 联系技术顾问：彭老师 13084473447（微信同号）")
            your_gemini_api_key = input("\n请重新输入您的Gemini API密钥: ").strip()
        
        # ================== 修改点2：移除默认密钥 ==================
        # 原问题代码：your_gemini_api_key = "*********" 
        # 安全隐患：公开代码包含有效密钥会被恶意滥用
        
        # ================== 修改点3：输入验证增强 ==================
        search_query, paper_count = get_user_input()
        while paper_count > 50:  # 添加合理限制
            print("\033[93m[注意] 单次检索数量超过50可能触发API限制\033[0m")
            paper_count = int(input("请输入合理检索数量 (1-50): "))
        
        # ================== 修改点4：动态输出路径 ==================
        output_file = input("请输入报告文件名（留空自动生成）: ").strip()
        if output_file and not output_file.endswith('.txt'):
            output_file += ".txt"
        
        # ================== 后续流程保持不变 ==================
        print(f"\n正在从PubMed检索关于「{search_query}」的{paper_count}篇文献...")
        pubmed_ids = search_pubmed(search_query, max_results=paper_count)
        if pubmed_ids:
            print(f"找到 {len(pubmed_ids)} 篇相关文献，PMID: {pubmed_ids}")
            
            # 显示部分摘要作为示例
            if pubmed_ids:
                first_abstract = fetch_abstracts([pubmed_ids[0]])
                print("\n第一篇文献摘要 (部分):\n", first_abstract[:500] + "...")

            print("\n开始处理各篇论文...")
            # 处理每篇论文 - 使用改进版函数
            paper_summaries = process_individual_papers_with_quota_management(pubmed_ids, your_gemini_api_key, search_query)
            
            print("\n获取所有论文原始摘要...")
            # 获取所有原始摘要用于综合总结
            all_abstracts = fetch_abstracts(pubmed_ids)
            
            print("\n生成综合总结...")
            # 生成综合总结
            overall_summary = summarize_abstracts_with_gemini(all_abstracts, your_gemini_api_key, search_query)
            
            print("\n创建并保存最终报告...")
            # 创建并保存报告
            report = create_literature_report(search_query, pubmed_ids, paper_summaries, overall_summary, output_file)
            
            print("\n报告生成完成！")
            print("\n按任意键退出程序...")
            input()

        else:
            print(f"未找到关于「{search_query}」的相关 PubMed 文献。请尝试使用不同的关键词。")
            print("\n按任意键退出程序...")
            input()
        
    except Exception as e:
        # ================== 修改点5：友好错误提示 ==================
        error_msg = f"""
        \n\033[91m[系统错误] {str(e)}\033[0m
        常见问题解决方案：
        1. API密钥无效 → 重新获取并输入有效密钥
        2. 网络连接失败 → 检查VPN/代理设置,本研究使用的是gemini apikey 建议老师配置合理的代理服务器访问海外网站
        3. 结果超限 → 减少检索数量
        4. 仍无法解决 → 截图此信息联系彭老师
        """
        print(error_msg)
        input("\n按回车键退出程序...")
